import logging
import functools
import os
import sys
import time
import matplotlib.pyplot as plt
import numpy as np
import re
import json
from tqdm import tqdm

logging.getLogger('matplotlib').setLevel(logging.CRITICAL)
def compute_ema(losses, alpha=0.1):
    ema_values = [losses[0]]  # start off with the first loss value

    for loss in losses[1:]:
        ema_values.append((1 - alpha) * ema_values[-1] + alpha * loss)

    return ema_values


def extract_tensor_info(file_path):
    """
    Read a text file and extract 'Tensor(shape=[], dtype=float32, value=0.xxxxx)'
    from each line until a stopping condition is met.

    Args:
        file_path (str): Path to the text file to be read.

    Returns:
        tensor_info_dict (dict): Dictionary with epoch numbers as keys and lists of
                                extracted tensor values as values.
    """
    tensor_infos = []  # Dictionary to store epoch-wise tensor values

    # Regular expression pattern to match the desired tensor value format
    tensor_value_pattern = r"Tensor\(shape=\[\], dtype=Float32, value=(0\.\d+)\)"

    # Stopping condition pattern with the updated timestamp format
    stopping_pattern = r"\[(\d{4}/\d{2}/\d{2} \d{2}:\d{2}:\d{2})\] - Epoch: (\d+), Time: \d+\.\d+, Steps: \d+\.\d+"

    # Open and read the file
    with open(file_path, 'r') as file:
        current_epoch = None
        tensor_values = []  # List to store tensor values for the current epoch

        for line in file:
            tensor_value_match = re.search(tensor_value_pattern, line)
            stopping_match = re.search(stopping_pattern, line)

            if tensor_value_match:
                tensor_value = float(tensor_value_match.group(1))
                tensor_values.append(tensor_value)

            if stopping_match:
                current_epoch = stopping_match.group(2)
                if current_epoch is not None:
                    tensor_infos.append(np.array(tensor_values))
                tensor_values = []

    # return tensor_infos
    if len(tensor_values) > 0:
        tensor_infos.append(np.array(tensor_values))
    return tensor_infos


def extract_loss_info(file_path):
    tensor_value_pattern = r"Tensor\(shape=\[\], dtype=Float32, value=(0\.\d+)\)"
    tensor_values = []  # List to store tensor values for the current epoch
    with open(file_path, 'r') as file:

        for line in file:
            tensor_value_match = re.search(tensor_value_pattern, line)

            if tensor_value_match:
                tensor_value = float(tensor_value_match.group(1))
                tensor_values.append(tensor_value)
    return np.array(tensor_values)


def cleanuptensors(tens):
    ll = []
    for t in tens:
        ll.append(len(t))

    tm = int(np.amin(ll))
    return [t[0:tm] for t in tens], tm



if __name__ == '__main__':
    # file_path = '/home/liupan/output/mm_newsize/print_rank_0.log'
    # ROOT = '/data/output/mm_vitbaseall3d2dmodeall2nodeslr3mskr060'
    # ROOT = '/storage/mindspore_env/output/mm_dev10nodesfullsizevitlarge'

    ROOT = str(sys.argv[1])

    F0 = os.path.join(ROOT, 'print_rank_0.log')
    print(f'extracting info from {F0}')
    tsinfos = extract_tensor_info(F0)
    nperepoch = len(tsinfos[0])
    print(f'training has reached {len(tsinfos)} epochs, and {nperepoch} per epoch')

    FA = [os.path.join(ROOT, i) for i in os.listdir(ROOT) if 'print_rank' in i and '_0' not in i]
    totaltensors = []
    for f in FA:
        totaltensors.append(extract_loss_info(f))

    totaltensors, tmin = cleanuptensors(totaltensors)
    tsinfos[-1] = np.array(tsinfos[-1][0:tmin-nperepoch*(len(tsinfos)-1)])
    total_ranks = len(totaltensors) + 1
    print(f'total in {total_ranks} ranks')

    newts = []
    for i, tt in tqdm(enumerate(tsinfos)):
        for ii in totaltensors:
            tt += ii[i * nperepoch:(i + 1) * nperepoch]
        ttt = tt / total_ranks
        newts.append(list(ttt))

    losses = []
    emaloss = []
    elosses = []
    eemalosses = []
    for i in range(len(newts)):
        _loss = newts[i]
        _emaloss = compute_ema(_loss)
        losses += _loss
        emaloss += _emaloss
        elosses.append(np.mean(_loss))
        eemalosses.append(np.mean(_emaloss))

    plt.plot(losses, label='loss')
    plt.plot(emaloss, label='loss_ema')
    plt.legend()
    plt.title('loss_step')
    plt.savefig(os.path.join(ROOT, f'loss_step.png'))
    plt.close()

    # plt.ylim([0, 0.1])
    plt.plot(elosses, label='loss')
    plt.plot(eemalosses, label='loss_ema')
    plt.legend()
    plt.title('loss_epoch')
    plt.savefig(os.path.join(ROOT, f'AEpoch{len(newts)}_loss_epoch.png'))
    plt.close()
